[WS1][CUDA][Qwen3-Next] C1: zero-centred RMSNorm references and CUDA kernel (RFC #428 stack 2/4) - #467

Local 2× B200, same software as above; extension rebuilt from source with
-gencode=arch=compute_100,code=sm_100. Measured at 5f5654a against main at
32b765e.

full suite, main (32b765e): 19 failed, 2940 passed, 923 skipped
full suite, 2de532c:        17 failed, 3013 passed, 923 skipped   (no new failures; 2940 + 2 + 6 + 65)

Each of the following was compared against main, on Python 3.10 with the hook versions
pinned in .pre-commit-config.yaml (CI pins Python 3.10):

• real pre-commit: no new findings;
• mypy --ignore-missing-imports rl_engine/: no new errors;
• mkdocs build --strict: 8 warnings, the same 8 as main.

At d03c852, pre-commit (same pinned hooks, Python 3.10) passes on the six files it
changes. mypy and mkdocs build --strict were not repeated there; d03c852 changes the
operator page, the module docstring, two comments and the tests' tolerance source.

The cluster numbers above are from one cluster run at d03c852. It built the extension in
place (TORCH_CUDA_ARCH_LIST=10.0 RL_KERNEL_REQUIRE_EXT=1 python setup.py build_ext --inplace) and checked that rl_engine._C loaded from that snapshot before running
anything. It ran check_operator without the shape flags; the values in the Test
section are the script's defaults (scripts/check_operator.py:107-120).

• Repository file:line refer to this branch's tip d03c852.
• vLLM paths are relative to the installed vllm 0.30.0 package root, e.g.
  model_executor/models/qwen3_next.py. transformers paths are relative to
  transformers/models/qwen3_next/ (5.17.0). After the first full path, files are cited
  by name.

vLLM imports the decoder norm as GemmaRMSNorm as Qwen3NextRMSNorm (qwen3_next.py:28).
Its forward_native does weight = self.weight.float() + 1.0
(model_executor/layers/layernorm.py:157), and its forward_cuda simply delegates to
forward_native (:163-168). So the +1 is an fp32 operation.

Not covered here:

• vLLM can route q_norm/k_norm through the fused Triton kernel
  fused_qk_rmsnorm_rope_gate (qwen3_next.py:405). That requires attn_output_gate,
  CUDA, neox-style RoPE, a supported dtype, and text_only or supports_mrope
  (:378-384). It is a separate provider, and where it applies the +1 has not been
  checked.
• vLLM reads the gated norm's activation from config.output_gate_type (default silu,
  sigmoid also allowed; qwen_gdn_linear_attn.py:480-485). The references here, like
  transformers, hard-code silu, which matches the released config.

RFC #428 §11.1 lists gated norm as a first-divergence trace boundary. The two upstreams
round it differently:

test_gated_conventions_diverge_in_low_precision asserts the consequence. In bf16 the two
conventions differ by more than 1e-3, and with fp32 inputs they agree bitwise, because
the cast is then a no-op. One earlier measurement had no committed script: one seed,
B200, bf16, head_v_dim=128, 512 rows, direct (eager) calls. It recorded:

HF          vs vLLM forward_native       same= 42617/65536 (65%)  max|d|=6.25e-2
HF          vs vLLM forward_cuda (FLA)   same= 42619/65536 (65%)  max|d|=6.25e-2

Both conventions are implemented. Qwen3NextRMSNormGatedOp follows vLLM: the weight
multiply is in fp32. Qwen3NextRMSNormGatedHFOp follows transformers, and is kept as a
named, tested witness so the divergence stays visible.

The convention matches, not the bits. All committed comparisons call forward_native and
forward_cuda directly, which is eager. Committed eager comparison over 40 seeds (bf16,
head_v_dim=128, 512 rows, one shape, B200):

ours vs forward_native            6/40 seeds differ, worst 1.56e-2
ours vs forward_cuda             18/40 seeds differ, worst 3.91e-3
forward_native vs forward_cuda   21/40 seeds differ, worst 1.56e-2

In fp32 the two eager paths differ on about 36% of elements. The test bounds only the
magnitude, max_abs <= 1e-5 (tests/check_qwen3_next_norm_providers.py:156), not a
per-element ULP bound.

vLLM's default mode compiles with inductor. custom_ops then gets "none"
(config/vllm.py:1608-1613), so the gated norm's forward_native is traced into the
model graph and compiled there (model_executor/custom_op.py:185-194, :296-311). The
gated norm is called outside the opaque GDN attention core (qwen_gdn_linear_attn.py:856,
:974), so it does reach the inductor graph.

What tests/check_qwen3_next_norm_providers.py pins:

• test_vllm_paths_disagree_only_within_bounds and
  test_ours_tracks_each_vllm_path_within_bounds bound the gap instead of asserting
  equality. The bounds are provider-gap bounds, not contract thresholds. They fail if
  either eager path's arithmetic drifts past the bounds. They cannot detect which path
  vLLM dispatches, for example the compiled one.
• test_gdn_decode_provider_env_defaults_are_recorded asserts the env defaults
  VLLM_GDN_DECODE_KERNEL="cuda" and VLLM_ENABLE_FLA_PACKED_RECURRENT_DECODE=1. It does
  not assert which kernel runs. For Qwen3-Next the first default does not take
  effect: gqa_interleaved_layout=True (qwen3_next.py:495) makes vLLM fall back to the
  Triton decode kernel.
• test_custom_op_has_distinct_native_and_cuda_paths asserts only that the two methods
  differ.

These are one-off probes, not committed checks and not vLLM engine runs. They are a
standalone inductor compile of the lowered implementation (IR torch-wrapping off), plus
vLLM's batch-invariant (BI) mode enabled directly. The scripts are not committed.

Setup:

• B200, bf16, 512 rows, seeds 0–39;
• weights: randn for gated, randn×0.1 for the zero-centred decoder;
• one cluster run, one compile mode per process;
• run at 99a01e7 (then the top of the stack). The PyTorch norm references the probe
  uses (rl_engine/kernels/ops/pytorch/norm/) differ from 5f5654a only by an added gate
  dtype/device check, so their arithmetic is unchanged. The probe does not use the CUDA
  norm code, which the gated PR changes substantially.

Validity checks: every process's dynamo graph count matched its compile mode (static 9,
dim-0 backed 2, dim-0 unbacked 1), and eager outputs hashed identically across processes.

Seeds differing out of 40:

Maximum |diff| is at most 1.56e-2 in every row: 3.91e-3 for the gated reference vs
compiled, and 0 for compiled with BI vs compiled without it.

What this shows:

• Compiling is a third numerical path for both norms. An earlier probe
  reported "compiled = eager" for the decoder. That result is withdrawn: IR torch-wrapping
  kept rms_norm opaque, so its mean was never compiled.
• In this standalone compile, BI mode does nothing. The batch-invariant mean is
  called 0 times from compiled code, and compiled output is identical with BI on or off.
  For a vLLM engine the expectation is the same, but untested.
• In eager, BI changes the decoder norm but not the gated norm's forward_native.
• For the decoder, compiled-under-BI and eager-under-BI agree on only 2–4 of 40 seeds, so
  the compiled path is not running BI's arithmetic.
• An earlier one-off probe also compared the gated norm's compiled path with
  FLA forward_cuda: 10/40 seeds differ, max 1.95e-3. That probe used torch.compile's
  default dynamic-shape setting, which differs from all three modes above.

Prefix-slice batch invariance (seeds 0–7, n ∈ {1,2,8,16,32,48,64,256} of 512):

• No failures for the dim-0 unbacked compile, default eager, or eager under BI, in
  all three cases.
• The static compile, which recompiles per shape, fails for the gated norm at
  n = 16/64/256, and for the residual case at n = 1.
• The dim-0 backed compile fails for the residual case at n = 1. That n = 1 graph is
  specialised separately, so it is not comparable to an engine.

The unbacked result resembles an engine only when there is a single compile range and no
compile_sizes. Passes such as fuse_allreduce_rms, enable_sp and the qk-norm/rope
fusions add range endpoints (config/vllm.py:2427-2525), and that behaviour is closer to
the static case. Whether they are on for Qwen3-Next TP1 by default was not checked.
Non-prefix slices (the RFC §2.1 "order" case) were not probed.

The next PR's C3 (forward) and C4 (gradient) gates cover the zero-centred CUDA op under
batch-size, chunk, padding and permutation transforms: every comparison is bitwise
(max_abs=0), hidden 2048, head-dim 128, a cluster run at 774ef70 on the
gated-norm branch. 774ef70 contains 5f5654a but not d03c852; d03c852 changes
only comments in csrc/cuda/rmsnorm.cu, so the kernel measured is the one in this PR.

shape_invariant_rstd (rl_engine/kernels/ops/pytorch/norm/rms_norm.py:100-118) came
from 8ed1693, which is on main (Zhang Jian, "fix(ascend): shape-invariant RMSNorm
rstd via a fixed-order reduction").

• That commit was motivated by NPU, where 52 of 140 rows flipped between [1,7,H] and
  [1,20,H] and the selected logp moved by 0.07.
• It was introduced as the shared reference for native, canonical and Ascend, and is
  documented as "bitwise identical for every layout on every device".

This PR adds one CUDA data point. Over 20 seeds at H=2048 in bf16, with one slicing
layout (x[3:5] and x[:1] of 64 rows), a plain mean(-1) broke slice invariance on
1 of 20 seeds; the chunked reduction broke on none.

• The test asserts only the chunked side (tests/test_qwen3_next_norm.py:184-200). By
  design it does not pin mean(-1) failing, so 1/20 is a recorded observation, not an
  assertion.
• The chunked reduction applies only when H % 32 == 0. Otherwise it falls back to
  mean(-1) (rms_norm.py:111-113). Qwen3-Next's 2048 and 128 both qualify.

This applies to the PyTorch references running on CUDA. The CUDA kernel has its own
fixed-order reduction: per-thread strided partial sums, then block_reduce_sum
(csrc/cuda/rmsnorm.cu:122-133). It is not bitwise equal to the reference; the test
compares the two under the contract's forward_accuracy reduction bf16 row,
atol=5e-2, rtol=2e-2 (tests/test_qwen3_next_norm.py:428-436).

RMSNormCudaOp.parameter_vjp_contributions_fp32 used to recompute rstd with
torch.rsqrt(x32.square().mean(dim=-1) + eps). It now takes rstd from
_C.rmsnorm_forward (rmsnorm.py:150-155), the same statistic the kernel's backward
uses. That costs one extra forward launch, whose y is discarded.

• Its only consumer is the gtest gradient adapter (gradient_adapters.py:820), reached
  through make_gradient_runner. In other words it is used by the C4 gate
  scripts/check_gradient_invariance.py, not by check_operator.
• test_cuda_parameter_contributions_reproduce_backward (8 cases, including
  offset=0.0) pins it.
• tests/test_ws1_gtest_gpu.py::test_c4_cuda_rms_norm_is_bitwise_invariant (:97) runs
  that C4 gate for the plain op. Its 9 passes above include this change.
• Whether the old implementation fails that equality was not run.

The decoder norm is not bitwise equal to any vLLM path probed (36–39 of 40 seeds
differ, max 1.56e-2; see the diagnostic probes above).

• Only layer 0's input norm takes the no-residual path (qwen3_next.py:559). The other 96
  of the model's 97 decoder-norm calls carry a residual (:561, :596, :729). They go
  through fused_add_rms_norm, which normalises the fp32 sum, and this PR has no
  reference for that path.
• The difference is attributed to reduction order, which has not been isolated. This PR
  uses a chunked fp32 sum. vLLM uses torch mean(-1) in eager, an inductor-generated
  reduction when compiled, and mean_batch_invariant in eager under BI.

An earlier version of this PR argued that matching vLLM bitwise would mean "trading L1 for
L2". That argument is withdrawn for lack of evidence. It rested on a single
observation that was never reproduced: mean(-1) broke slice invariance on 1 of 20
seeds, and which of the checked slices failed (x[3:5] or x[:1]) was not recorded.
Every vLLM path passed the prefix-slice probe above, but non-prefix slices were not
probed.

About RFC §0 item 3 / §5.2 and VLLM_BATCH_INVARIANT=1: it replaces aten::mean.dim
with mean_batch_invariant (model_executor/determinism/batch_invariant.py:900,
:1119) on aten's CUDA dispatch key, so it acts only where mean is dispatched eagerly.

• Decoder norm: in enforce_eager the decoder norm's forward_cuda delegates to
  forward_native (layernorm.py:163-168), whose mean (ir/ops/layernorm.py:17) BI
  does replace. Of these two norms, this is the only case where the flag changes the
  arithmetic. Elsewhere BI also changes matmul, the plain RMSNorm
  (layernorm.py:101-112) and softmax.
• Gated norm: enforce_eager sets custom_ops to "all"
  (config/vllm.py:1608-1615), so RMSNormGated runs forward_cuda, i.e. the FLA
  Triton kernel. That never reaches aten.mean, so BI cannot affect it. The relevant
  comparison there is this PR's reference vs forward_cuda (18/40 seeds differ).
• Default compiled mode: the standalone probe shows the compiled graph bypasses BI's
  mean.

• RFC §7.2 lists RMSNorm/residual as "vLLM strict forward plus custom backward | Reuse",
  and §9.1's C1 row names "vLLM strict ops" as the provider.
• This PR instead adds an independent PyTorch reference and the repo's own CUDA kernel.
• One reading of "vLLM strict forward" is vLLM under VLLM_BATCH_INVARIANT=1 (see
  above). The RFC itself cautions that "setting the environment variable alone is not
  acceptance evidence" (§5.2).
• The existing Qwen3 dense strict profile in this repo does the opposite of the RFC plan.
  _patch_qwen3_strict_model (rl_engine/integrations/vllm_runtime.py:1049, "Align
  vLLM's RMSNorm and Attention projections with Megatron") replaces vLLM's RMSNorm with
  _strict_rms_norm, which "Preserve[s] PyTorch eager RMSNorm arithmetic"
  (rms_norm.py:10-15).
• Nothing in rl_engine/integrations/ on this branch patches any Qwen3-Next norm.

Which source of truth should the Qwen3-Next CUDA strict profile use?

• (a) The RFC plan: a shared vLLM forward under VLLM_BATCH_INVARIANT=1. This PR's
  references then become test oracles for it. What (a) requires is that training
  reproduce the arithmetic rollout actually executes. As the probes above show, BI alters
  that arithmetic only for the decoder norm, and only in enforce_eager. In the default
  compiled mode the flag has no effect on these norms. There, batch invariance is whatever
  inductor's generated reduction provides, which was tested only for prefix slices with a
  single compile range.
• (b) The Qwen3 dense precedent: vLLM patched to the repo's arithmetic.
• (c) The repo's own CUDA kernels on both sides.

The answer also decides which gated convention is the strict default.
docs/design/ws1-numerical-precision-standard.md §11.1 already lists "declare the RMSNorm
rounding order ... and its relationship to the training-side framework" as
release-blocking.

Earlier commits on this branch had put statements in the operator page, the module
docstring and the provider check that this description does not make. d03c852 corrects
them in-tree. The main corrections are below; the full list is in its commit message.

• the operator page no longer says the op is registered or prints a check_operator
  command (registration arrives with the gated-norm PR);
• withdrawn: "7 of 1048576 differ" (single seed, no script), the 35% / 5.3e-2 cast-order
  isolation (an fp32 round-trip is a no-op), "needed on CUDA as well", and "matching vLLM
  means trading L1 for L2" (one unreproduced observation);
• the provider check no longer claims to establish which path vLLM dispatches, or a
  per-element fp32 ULP bound;
• tests resolve tolerances from the contract (reduction row) instead of the elementwise
  bf16 numbers. That loosens bf16 (5e-2/2e-2) and tightens fp16 (1e-3/1e-3, previously
  2e-2); the fp16 cases still pass.

History on this branch is not rewritten, so 6741d16's commit message still carries the
withdrawn statements, plus an unscoped "L0 repeatable and L1 batch-invariant" claim
level (the scoped levels are in the Summary), "default offset vs the previous kernel —
bitwise equal" (no test compares against the previous kernel) and "a vLLM bump that
changes the provider fails loudly" (see the narrower statement above). Where they
conflict, this description and d03c852 are the corrected versions.

main's CI history as of 2026-10-01 (GitHub Actions run list):

• CI-Pipeline fails at both 32b765e and the current 1968a87: linting fails at
  Run pre-commit hooks, and unit-tests is then skipped. Docs fails at both commits
  as well.
• WS1-chain-GPU fails at both commits at Configure runpodctl, before any test runs.
  ws1-gtest-gpu failed at the same step on each of its last five push-to-main runs
  (latest b31fac9, 2026-09-21).
• The cause is step order, not a recent release. runpodctl config --apiKey, run as the
  first runpodctl call on a fresh runner, fails because no config file exists yet. Run
  36835086246 logs 'runpodctl config' is deprecated followed by error saving config,
  which reproduces locally with v2.14.0 on an empty HOME. ws1-gtest-gpu has never
  succeeded: as of 2026-10-01, 122 runs (69 failed, 40 skipped, 13 action_required),
  0 successes, and its first run (2026-08-19) already
  failed at this step.
• gpu-ci.yml:50 uses the same command but runs runpodctl version first, which creates
  the config file, so it got past that step on its last run that reached RunPod
  (2026-09-11, run 34585491929). That run then failed later, at pod create. A separate
  fix for the setup step is being prepared. No gpu-ci run has executed since 2026-09-11,
  so whether it still gets past setup today is unverified.
• ws1-chain-npu was queued on 1968a87 and cancelled on 32b765e.

As of today, no CI path would execute these tests, before or after merge.
ws1-gtest-gpu and ws1-chain-gpu are triggered on push to main but stop at RunPod
setup. gpu-ci is triggered only by pull_request_target (gpu-ci.yml:4), never by a
merge; it last stopped at pod create, and its pytest tests/ would in any case stop at
the tests/test_rocm_aiter_api_contract.py collection error before running anything.

If those problems are fixed, the workflows' scripts would run:

1. ci.yml: unit-tests has needs: linting. Once green, it runs this file on a
   CPU-only torch (ci.yml:63): the 38 CPU cases run, and the 27 CUDA cases skip.
2. ws1-gtest-gpu.yml: this PR's changes under rl_engine/kernels/ops/** match its
   paths. It runs for same-repo PRs and on push to main, on sm_86 and sm_90
   runners, and is skipped for fork PRs. Its script ci/run_ws1_gtest.sh does not run
   tests/test_qwen3_next_norm.py. It does run tests/test_ws1_gtest_gpu.py, whose C4
   gate exercises the plain op's parameter-VJP change; its C8 sweep also constructs
   RMSNormCudaOp directly.
3. ws1-chain-gpu.yml: runs on push and PR to main/test, with no paths
   filter and no needs:, and is skipped for fork PRs (:46). It runs
   ci/run_ws1_chain_gate.sh: three pytest files (test_kv_consistency.py,
   test_ws1_qwen3_dense.py, test_ws1_chain_integration.py) and the C8 sweep. That
   chain constructs RMSNormCudaOp directly from its gtest candidate path
   (rl_engine/alignment/qwen3_dense.py:374-396), not through get_op.
4. gpu-ci.yml (needs the needs-gpu-ci label; sm_86/sm_90; pins
   transformers==5.13.1): runs pytest tests/ on a GPU. With the collection error
   above fixed, it would collect the 27 CUDA cases.
5. The check_ provider file needs vLLM and runs in none of these.

For a PR opened from a fork, (4) is the only route before merge, and only after a
maintainer adds the label.

• These are sm_100 numbers, not the sm_90 the project's benchmarks usually use. Per RFC
  §2.2, no claim is made across H100/B200.
• Stacked on fix/registry-cuda-rmsnorm-dispatch, which must merge first; this branch
  contains its commits. d03c852 merges cleanly into the current upstream/test-qwennext
  (11cac8c, git merge-tree --write-tree); 32b765e..11cac8c is formatting and lint/docs/CI
  configuration plus the logp indexed-write fix, none of which changes the arithmetic measured
  here, so the main measurements above carry over. Two further PRs build
  on this one.

• tests/test_qwen3_next_norm.py defines its own _rand and requires_cuda_rmsnorm
  (:49, :362), which duplicate helpers in tests/test_rms_norm.py (:34, :109).
• The GDN gated norm has no CUDA kernel in this PR; that is the next PR. The gated
  PyTorch pair is not documented on this branch beyond a pointer: the operator page
  says it is documented with its CUDA kernel in the gated-norm PR
  (docs/operators/qwen3-next-rms-norm.md:114-115).
• 5f5654a has a subject-only message; its rationale is in the commit table above.

None of these changes behavior. They are left for follow-ups, because any commit here
would have to be merged forward into the two branches stacked on it.